# optimal_knn_webapp_pinecone.py
import gradio as gr
import cv2
import numpy as np
from PIL import Image
from pinecone import Pinecone

# ---------- 初始化 Pinecone ----------
pc = Pinecone(api_key="pcsk_4RZEQz_7gy8B1tBbXWFw22UX6swnKLafVduhYCrTx4eULaxrpgsWEEAnt96Y7hXstdgFqi")
index = pc.Index("mnist-index")

def predict_digit(inp):
    if not isinstance(inp, dict):
        return f"输入不是 dict，实际类型：{type(inp)}"

    if "composite" in inp:
        img = inp["composite"]
    elif "image" in inp:
        img = inp["image"]
    elif "layers" in inp:
        img = inp["layers"][0] if len(inp["layers"]) else None
    else:
        return "找不到可用图字段"

    gray = cv2.cvtColor(np.array(img), cv2.COLOR_RGB2GRAY)
    gray = cv2.bitwise_not(gray)  # 白底黑字 → 黑底白字
    resized = cv2.resize(gray, (8, 8), interpolation=cv2.INTER_AREA)

    # 调试：打印 8×8 图像矩阵
    print("调试：8x8 图像矩阵（放大 16 倍便于查看）：")
    print(np.round(resized / 16.0 * 16).reshape(8, 8))

    vector = (resized / 16.0).reshape(-1).tolist()

    result = index.query(vector=vector, top_k=11, include_metadata=True)
    if not result.matches:
        return "未识别"
    labels = [int(match.metadata["label"]) for match in result.matches]
    pred = max(set(labels), key=labels.count)
    return str(pred)

# ---------- Gradio 界面 ----------
with gr.Blocks(title="Pinecone KNN 手写数字识别") as demo:
    gr.Markdown("### 基于 Pinecone 的 KNN 手写数字识别")
    with gr.Row():
        pad = gr.Sketchpad(height=200, width=200, image_mode="RGB")
        lbl = gr.Label()
    btn = gr.Button("Submit")
    btn.click(predict_digit, inputs=pad, outputs=lbl)

demo.launch()